import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import keras.backend as K
import tensorflow as tf
from tensorflow.python.client import device_lib
import keras
def run():
    print()
    b = input("是否显示本地CPU/GPU/TPU的详细信息？True or False")
    assert isinstance(b,bool),TypeError
    if b:
        print('本地CPU/GPU/TPU的详细信息：\n',device_lib.list_local_devices())
    print(f"tensorflow 的版本为： {tf.__version__}")
    print(f"tensorflow 的版本为： {keras.__version__}")
    print("开始检测是否安装好cuda...")
    if tf.test.is_built_with_cuda():
        print("cuda安装完好")
        print('查看可用GPU...')
        gpu_lst = K.tensorflow_backend._get_available_gpus()
        n_gou = len(gpu_lst)
        print(f"一共{n_gou}个GPU可以使用，列表如下：\n{gpu_lst}")
        if tf.test.gpu_device_name():
            print(f'默认使用设备为：{tf.test.gpu_device_name()}')
        else:
            print("Please install GPU version of TF")
    else:
        print("cuda未安装")

if __name__ == '__main__':
    run()
    print(device_lib.list_local_devices())

